/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */


#include <faiss/impl/pq4_fast_scan.h>

#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/simd_result_handlers.h>


namespace faiss {


using namespace simd_result_handlers;

/***************************************************************
 * accumulation functions
 ***************************************************************/

namespace {

/*
 * The computation kernel
 * It accumulates results for NQ queries and BB * 32 database elements
 * writes results in a ResultHandler
 */

template<int NQ, int BB, class ResultHandler>
void kernel_accumulate_block(
        int nsq,
        const uint8_t *codes,
        const uint8_t *LUT,
        ResultHandler & res)
{
    // distance accumulators
    simd16uint16 accu[NQ][BB][4];

    for(int q = 0; q < NQ; q++) {
        for(int b = 0; b < BB; b++) {
            accu[q][b][0].clear();
            accu[q][b][1].clear();
            accu[q][b][2].clear();
            accu[q][b][3].clear();
        }
    }

    for(int sq = 0; sq < nsq; sq += 2) {
        simd32uint8 lut_cache[NQ];
        for(int q = 0; q < NQ; q++) {
            lut_cache[q] = simd32uint8(LUT);
            LUT += 32;
        }

        for (int b = 0; b < BB; b++) {
            simd32uint8 c = simd32uint8(codes);
            codes += 32;
            simd32uint8 mask(15);
            simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
            simd32uint8 clo = c & mask;

            for(int q = 0; q < NQ; q++) {
                simd32uint8 lut = lut_cache[q];
                simd32uint8 res0 = lut.lookup_2_lanes(clo);
                simd32uint8 res1 = lut.lookup_2_lanes(chi);

                accu[q][b][0] += simd16uint16(res0);
                accu[q][b][1] += simd16uint16(res0) >> 8;

                accu[q][b][2] += simd16uint16(res1);
                accu[q][b][3] += simd16uint16(res1) >> 8;
            }
        }
    }

    for(int q = 0; q < NQ; q++) {
        for (int b = 0; b < BB; b++) {

            accu[q][b][0] -= accu[q][b][1] << 8;
            simd16uint16 dis0 = combine2x2(accu[q][b][0], accu[q][b][1]);

            accu[q][b][2] -= accu[q][b][3] << 8;
            simd16uint16 dis1 = combine2x2(accu[q][b][2], accu[q][b][3]);

            res.handle(q, b, dis0, dis1);
        }
    }


}


template<int NQ, int BB, class ResultHandler>
void accumulate_fixed_blocks(
        size_t nb,
        int nsq,
        const uint8_t *codes,
        const uint8_t *LUT,
        ResultHandler & res)
{
    constexpr int bbs = 32 * BB;
    for (int64_t j0 = 0; j0 < nb; j0 += bbs) {
        FixedStorageHandler<NQ, 2 * BB> res2;
        kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2);
        res.set_block_origin(0, j0);
        res2.to_other_handler(res);
        codes += bbs * nsq / 2;
    }
}


} // anonymous namespace

template<class ResultHandler>
void pq4_accumulate_loop(
        int nq,
        size_t nb, int bbs,
        int nsq,
        const uint8_t *codes,
        const uint8_t *LUT,
        ResultHandler & res)
{
    FAISS_THROW_IF_NOT(is_aligned_pointer(codes));
    FAISS_THROW_IF_NOT(is_aligned_pointer(LUT));
    FAISS_THROW_IF_NOT(bbs % 32 == 0);
    FAISS_THROW_IF_NOT(nb % bbs == 0);

#define DISPATCH(NQ, BB) \
   case NQ * 1000 + BB: \
    accumulate_fixed_blocks<NQ, BB>(nb, nsq, codes, LUT, res);  \
    break

    switch(nq * 1000 + bbs / 32) {
        DISPATCH(1, 1);
        DISPATCH(1, 2);
        DISPATCH(1, 3);
        DISPATCH(1, 4);
        DISPATCH(1, 5);
        DISPATCH(2, 1);
        DISPATCH(2, 2);
        DISPATCH(3, 1);
        DISPATCH(4, 1);
    default:
        FAISS_THROW_FMT("nq=%d bbs=%d not instantiated", nq, bbs);
    }
#undef DISPATCH

}

// explicit template instantiations




#define INSTANTIATE_ACCUMULATE(TH, C, with_id_map) \
template void pq4_accumulate_loop<TH<C, with_id_map>> \
    (int, size_t, int, int, const uint8_t *, const uint8_t *, TH<C, with_id_map> &);

#define INSTANTIATE_3(C, with_id_map) \
INSTANTIATE_ACCUMULATE(SingleResultHandler, C, with_id_map) \
INSTANTIATE_ACCUMULATE(HeapHandler, C, with_id_map) \
INSTANTIATE_ACCUMULATE(ReservoirHandler, C, with_id_map) \

using Csi = CMax<uint16_t, int>;
INSTANTIATE_3(Csi, false);
using CsiMin = CMin<uint16_t, int>;
INSTANTIATE_3(CsiMin, false);

using Csl = CMax<uint16_t, int64_t>;
INSTANTIATE_3(Csl, true);
using CslMin = CMin<uint16_t, int64_t>;
INSTANTIATE_3(CslMin, true);



} // namespace faiss

